import numpy as np
import pandas as pd
import torch


# predict_utils.py
def add_time_features(df):
    df['hour_sin'] = np.sin(2 * np.pi * df.index.hour / 24)
    df['hour_cos'] = np.cos(2 * np.pi * df.index.hour / 24)
    df['month_sin'] = np.sin(2 * np.pi * df.index.month / 12)
    df['month_cos'] = np.cos(2 * np.pi * df.index.month / 12)
    df['dayofweek_sin'] = np.sin(2 * np.pi * df.index.dayofweek / 7)
    df['dayofweek_cos'] = np.cos(2 * np.pi * df.index.dayofweek / 7)
    return df


def create_sequences(data, seq_length):
    """
    创建输入序列和目标序列（所有特征均为目标）
    Args:
        data: DataFrame，包含所有特征
        seq_length: 序列长度
    Returns:
        X: 输入序列 (num_samples, seq_length, num_features)
        y: 目标序列 (num_samples, num_features)
    """
    X, y = [], []
    for i in range(len(data) - seq_length):
        X.append(data.iloc[i:i + seq_length].values)
        y.append(data.iloc[i + seq_length].values)
    return np.array(X), np.array(y)


def predict_time_range(model, data_scaler, data, start_time, num_steps, seq_length=672):
    """
    输入时间点和预测步数，返回多特征预测结果
    Args:
        model: 训练好的模型
        data_scaler: 数据标准化器（StandardScaler）
        data: 原始数据 DataFrame（带时间索引）
        start_time: 起始时间点（字符串，如 '2025-01-15 23:30'）
        num_steps: 需要预测的时间步数
        seq_length: 模型训练时的序列长度
    Returns:
        pred_df: DataFrame，形状为 (num_steps, num_features)
    """
    # 转换为时间戳对象
    start_time = pd.to_datetime(start_time)
    max_train_time = data.index.max()
    allowed_max_time = max_train_time + pd.Timedelta(days=7)

    # 验证时间范围
    if start_time > allowed_max_time:
        raise ValueError(f"起始时间 {start_time} 超出允许范围。允许的最晚时间为 {allowed_max_time}")

    # 获取历史数据
    if start_time in data.index:
        # 常规预测模式
        start_idx = data.index.get_loc(start_time)
        historical_data = data.iloc[start_idx - seq_length:start_idx]
    else:
        # 扩展预测模式：使用最后seq_length个时间步
        historical_data = data.iloc[-seq_length:]

    historical_scaled = data_scaler.transform(historical_data)

    # 递归预测
    predictions = []
    current_seq = historical_scaled.copy()

    for _ in range(num_steps):
        input_tensor = torch.tensor(current_seq[np.newaxis, :, :], dtype=torch.float32)
        with torch.no_grad():
            pred_scaled = model(input_tensor).numpy()
        pred_original = data_scaler.inverse_transform(pred_scaled)
        predictions.append(pred_original[0])
        new_data = data_scaler.transform(pred_original).reshape(1, -1)
        current_seq = np.vstack([current_seq[1:], new_data])

    # 构建结果DataFrame
    pred_df = pd.DataFrame(
        predictions,
        index=pd.date_range(start=start_time, periods=num_steps, freq='30min'),
        columns=data.columns
    )
    return pred_df
